package org.example.service;


import com.baomidou.mybatisplus.generator.FastAutoGenerator;
import com.baomidou.mybatisplus.generator.config.DataSourceConfig;
import com.baomidou.mybatisplus.generator.config.querys.DMQuery;
import com.baomidou.mybatisplus.generator.config.querys.KingbaseESQuery;
import com.baomidou.mybatisplus.generator.config.querys.MySqlQuery;
import com.baomidou.mybatisplus.generator.config.querys.PostgreSqlQuery;
import com.baomidou.mybatisplus.generator.config.rules.DateType;
import com.baomidou.mybatisplus.generator.engine.FreemarkerTemplateEngine;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Comparator;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

/**
 * 代码生成器控制器
 */
@RestController
public class GeneratorController {

    @PostMapping("/generate")
    public ResponseEntity<ByteArrayResource> generateCode(@RequestBody GeneratorRequest request) {
        Path tempDir = null;
        try {
            // 1. 创建临时目录用于存放生成的代码
            tempDir = Files.createTempDirectory("mybatis-gen-");
            String outputDir = tempDir.toAbsolutePath().toString();

            // 2. 配置数据源
            DataSourceConfig.Builder dataSourceBuilder;
            String jdbcUrl;

            switch (request.getDbType().toUpperCase()) {
                case "MYSQL":
                    jdbcUrl = String.format("jdbc:mysql://%s:%d/%s?useUnicode=true&characterEncoding=utf8&useSSL=false&serverTimezone=Asia/Shanghai",
                            request.getHost(), request.getPort(), request.getDatabase());
                    dataSourceBuilder = new DataSourceConfig.Builder(jdbcUrl, request.getUsername(), request.getPassword())
                            .dbQuery(new MySqlQuery()) // 使用 MySqlQuery
                            .schema(request.getDatabase()); // 设置 schema
                    break;
                case "DM":
                    // 达梦数据库URL示例，schema通常在URL参数中或通过setSchema
                    jdbcUrl = String.format("jdbc:dm://%s:%d/%s?schema=%s&characterEncoding=utf-8&autoReconnect=true",
                            request.getHost(), request.getPort(), request.getDatabase(), request.getDatabase()); // 达梦的database通常就是schema
                    dataSourceBuilder = new DataSourceConfig.Builder(jdbcUrl, request.getUsername(), request.getPassword())
                            .dbQuery(new DMQuery()) // 使用 DmQuery
                            .schema(request.getDatabase()); // 设置 schema
                    break;
                case "KINGBASEES":
                    // 人大金仓数据库URL示例
                    jdbcUrl = String.format("jdbc:kingbase8://%s:%d/%s",
                            request.getHost(), request.getPort(), request.getDatabase());
                    dataSourceBuilder = new DataSourceConfig.Builder(jdbcUrl, request.getUsername(), request.getPassword())
                            .dbQuery(new KingbaseESQuery()) // 使用 KingbaseESQuery
                            .schema(request.getDatabase()); // 设置 schema
                    break;
                case "POSTGRESQL":
                    jdbcUrl = String.format("jdbc:postgresql://%s:%d/%s",
                            request.getHost(), request.getPort(), request.getDatabase());
                    dataSourceBuilder = new DataSourceConfig.Builder(jdbcUrl, request.getUsername(), request.getPassword())
                            .dbQuery(new PostgreSqlQuery()) // 使用 PostgreSqlQuery
                            .schema(request.getDatabase()); // 设置 schema
                    break;
                default:
                    throw new IllegalArgumentException("不支持的数据库类型: " + request.getDbType());
            }

            // 3. 执行代码生成
            FastAutoGenerator.create(dataSourceBuilder)
                    .globalConfig(builder -> builder
                            .author(request.getAuthor()) // 作者
                            .outputDir(outputDir) // 输出目录
                            .commentDate("yyyy-MM-dd") // 注释日期格式
                            .dateType(DateType.TIME_PACK) // 日期类型，使用java.time
                            .enableSwagger() // 启用 Swagger 2
                            .disableOpenDir() // 生成后不打开目录
                    )
                    .packageConfig(builder -> builder
                            .parent(request.getParentPackage()) // 父包名
                            .entity("entity") // 实体类包名
                            .mapper("mapper") // Mapper 接口包名
                            .service("service") // Service 接口包名
                            .serviceImpl("service.impl") // Service 实现类包名
                            .xml("mapper.xml") // Mapper XML 包名
                    )
                    .strategyConfig(builder -> {
                        builder.entityBuilder().enableLombok(); // 在 Entity.Builder 上调用 enableLombok()
                        if (request.getTableNames() != null && !request.getTableNames().isEmpty()) {
                            // 修复：addInclude 应该在 StrategyConfig.Builder 上调用
                            builder.addInclude(request.getTableNames()); // 指定要生成的表
                        } else {
                            // 如果不指定表名，则默认生成所有表（或根据需要配置排除表）
                            // builder.addExclude("sys_log"); // 示例：排除某些表
                        }
                    })
                    .templateEngine(new FreemarkerTemplateEngine()) // 使用 Freemarker 模板引擎
                    .execute();

            // 4. 将生成的代码打包成 ZIP
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            try (ZipOutputStream zos = new ZipOutputStream(baos)) {
                zipDirectory(tempDir.toFile(), tempDir.toFile().getName(), zos);
            }

            // 5. 返回 ZIP 文件
            ByteArrayResource resource = new ByteArrayResource(baos.toByteArray());

            return ResponseEntity.ok()
                    .header(HttpHeaders.CONTENT_DISPOSITION, "attachment;filename=generated_code.zip")
                    .contentType(MediaType.APPLICATION_OCTET_STREAM)
                    .contentLength(resource.contentLength())
                    .body(resource);

        } catch (Exception e) {
            e.printStackTrace();
            return ResponseEntity.status(500).body(new ByteArrayResource(e.getMessage().getBytes()));
        } finally {
            // 6. 清理临时目录
            if (tempDir != null && Files.exists(tempDir)) {
                try {
                    Files.walk(tempDir)
                            .sorted(Comparator.reverseOrder())
                            .map(Path::toFile)
                            .forEach(File::delete);
                } catch (IOException e) {
                    System.err.println("无法删除临时目录: " + tempDir + ", 错误: " + e.getMessage());
                }
            }
        }
    }

    /**
     * 递归压缩目录
     * @param folder 要压缩的文件夹
     * @param parentEntryName ZIP文件中当前文件夹的路径名
     * @param zos ZipOutputStream
     * @throws IOException
     */
    private void zipDirectory(File folder, String parentEntryName, ZipOutputStream zos) throws IOException {
        for (File file : folder.listFiles()) {
            if (file.isDirectory()) {
                zipDirectory(file, parentEntryName + "/" + file.getName(), zos);
            } else {
                try (FileInputStream fis = new FileInputStream(file)) {
                    ZipEntry zipEntry = new ZipEntry(parentEntryName + "/" + file.getName());
                    zos.putNextEntry(zipEntry);
                    byte[] bytes = new byte[1024];
                    int length;
                    while ((length = fis.read(bytes)) >= 0) {
                        zos.write(bytes, 0, length);
                    }
                    zos.closeEntry();
                }
            }
        }
    }
}
